package com.yzh.zhipu.model;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yzh.zhipu.model.args.Prompt;
import com.yzh.zhipu.model.vo.Model;
import lombok.Data;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Data
@JsonInclude(JsonInclude.Include.NON_NULL)
public class ChatRequest {

    //---------------------v3 & v4版本通用参数 start---------------------
    private Model model = Model.GLM_3_TURBO;

    /**
     * 请求ID
     */
    @JsonProperty("request_id")
    private String requestId = String.format("yzh-%d", System.currentTimeMillis());
    /**
     * 采样温度，控制输出的随机性
     */
    private float temperature = 0.9f;
    /**
     * 温度取样的另一种方法，不要同时调整temperature 和 top_p
     */
    @JsonProperty("top_p")
    private float topP = 0.7f;
    //---------------------v3 & v4版本通用参数 end---------------------

    //======================= v4版本参数 start =======================
    /**
     * 输入给模型的会话信息
     * 用户输入的内容；role=user 挟带历史的内容；role=assistant
     */
    private List<Prompt> messages;

    /**
     * 启用采样策略（temperature&topP）开启或关闭
     */
    @JsonProperty("do_sample")
    private Boolean doSample = true;

    /**
     * true = sse请求，false = 同步请求
     */
    private Boolean stream = true;

    /**
     * 模型输出最大tokens
     */
    @JsonProperty("max_tokens")
    private Integer maxTokens = 2048;

    /**
     * 模型在遇到stop所制定的字符时将停止生成，目前仅支持单个停止词，格式为["stop_word1"]
     */
    private List<String> stop;
    //======================= v4版本参数 end =======================

    //----------------------- v3版本参数 start -----------------------
    /**
     * 输入给模型的会话信息
     * 用户输入的内容；role=user 挟带历史的内容；role=assistant
     */
    private List<Prompt> prompt;
    /**
     * 智普AI sse 固定参数 incremental = true 【增量返回】
     */
    private boolean incremental = true;
    /**
     * sseformat, 用于兼容解决sse增量模式okhttpsse截取data:后面空格问题, [data: hello]。只在增量模式下使用sseFormat。
     */
    private String sseFormat = "data";
    //----------------------- v3版本参数 end -----------------------


    @Override
    public String toString() {
        Map<String, Object> paramsMap = new HashMap<>();
        try {
            if (Model.CHATGLM_TURBO.equals(this.model)) {
                //v3旧模型走这里
                paramsMap.put("prompt", prompt);
                paramsMap.put("incremental", incremental);
                paramsMap.put("sseFormat", sseFormat);
            } else {
                //v4新模型走这里
                paramsMap.put("messages", messages);
                paramsMap.put("do_sample", doSample);
                paramsMap.put("stream", stream);
                paramsMap.put("max_tokens", maxTokens);
                paramsMap.put("stop", stop);
            }
            //v3 & v4 共有的参数
            paramsMap.put("model", model.getCode());
            paramsMap.put("request_id", requestId);
            paramsMap.put("temperature", temperature);
            paramsMap.put("top_p", topP);
            return new ObjectMapper().writeValueAsString(paramsMap);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }

}
